Skip to content

baml-language/fix: use build_plan for llm functions, delete execute#3279

Open
sxlijin wants to merge 1 commit intocanaryfrom
push-vmxnwtqstkwo
Open

baml-language/fix: use build_plan for llm functions, delete execute#3279
sxlijin wants to merge 1 commit intocanaryfrom
push-vmxnwtqstkwo

Conversation

@sxlijin
Copy link
Copy Markdown
Contributor

@sxlijin sxlijin commented Mar 25, 2026

Summary by CodeRabbit

  • Refactor

    • Restructured LLM call orchestration to use step-based execution with per-step delay handling and improved fallback management.
    • Enhanced error handling to provide clearer messaging when orchestration steps fail.
    • Improved type inference for client and retry policy definitions.
  • Tests

    • Updated orchestration integration tests to align with new step-based execution model.

@vercel
Copy link
Copy Markdown

vercel bot commented Mar 25, 2026

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
beps Ready Ready Preview, Comment Mar 25, 2026 9:53pm
promptfiddle Ready Ready Preview, Comment Mar 25, 2026 9:53pm

Request Review

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Mar 25, 2026

📝 Walkthrough

Walkthrough

The changes refactor the LLM execution flow from direct client execution to a step-based orchestration model. The ExecutionResult class is removed, and call_llm_function<T> now builds execution plans, advances round-robin state, iterates through steps with per-step delays, and returns the first successful result. Supporting type system and test infrastructure are updated accordingly.

Changes

Cohort / File(s) Summary
Core LLM Orchestration
baml_language/crates/baml_builtins2/baml_std/baml/llm.baml, baml_language/crates/baml_builtins2/baml_std/baml/llm_types.baml
Removed ExecutionResult class. Refactored call_llm_function to use build_plan(), advance_round_robin(), and per-step execute_step() with delays instead of single client.execute() call. Added advance_round_robin() method to Client and new top-level execute_step<T>() function; removed Client.execute<T>() and Client.execute_once<T>() methods.
Type Inference
baml_language/crates/baml_compiler2_tir/src/builder.rs
Added type inference for Definition::Let bindings in infer_single_name(), mapping LetOrigin::Client to Ty::Class(baml.llm.Client) and LetOrigin::RetryPolicy to Ty::Class(baml.llm.RetryPolicy).
Test Infrastructure
baml_language/crates/bex_engine/tests/orchestration.rs
Refactored orchestration tests to use direct client method invocation (A.build_plan(), FB.build_plan(), etc.) instead of module-level baml.llm.build_plan(). Removed helper function F(x: string), added new diagnostic test diag_fallback_fields, and standardized ignore attributes with updated diagnostic reasons.

Sequence Diagram(s)

sequenceDiagram
    participant Caller
    participant CallLLMFunc as call_llm_function
    participant Client
    participant Orchestration as Orchestration Step
    participant ErrorHandler

    Caller->>CallLLMFunc: call_llm_function<T>(context)
    CallLLMFunc->>Client: build_plan()
    Client-->>CallLLMFunc: steps: [OrchestrationStep]
    
    CallLLMFunc->>Client: advance_round_robin()
    Client-->>CallLLMFunc: state updated
    
    loop For each step in steps
        CallLLMFunc->>CallLLMFunc: check step.delay_ms
        alt delay_ms > 0
            CallLLMFunc->>CallLLMFunc: sleep(delay_ms)
        end
        
        CallLLMFunc->>Orchestration: execute_step(step, context)
        
        alt Step succeeds
            Orchestration-->>CallLLMFunc: result: T
            CallLLMFunc-->>Caller: return result
        else Step fails
            Orchestration-->>ErrorHandler: exception caught
            ErrorHandler-->>CallLLMFunc: continue to next step
        end
    end
    
    alt All steps failed
        CallLLMFunc-->>CallLLMFunc: throw DevOther("All orchestration steps failed")
        CallLLMFunc-->>Caller: error
    end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check ✅ Passed The PR title accurately describes the main change: refactoring LLM functions to use build_plan instead of the execute method, and removing the execute implementation.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch push-vmxnwtqstkwo

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 11ce9f11b9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

let result: T = client.execute(context, 0);
result
let steps = client.build_plan();
client.advance_round_robin();
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Advance round-robin counters only for attempted steps

call_llm_function now mutates round-robin state before any step runs, which breaks fallback behavior when an earlier branch succeeds. Because build_plan() flattens all fallback branches but the loop returns on the first successful step, many planned steps are never attempted; however client.advance_round_robin() still increments counters in those untouched subtrees, so future calls can skip providers that were never actually used. This is a regression from the previous execute-on-visit behavior and changes routing deterministically for Fallback[..., RoundRobin[...]] clients.

Useful? React with 👍 / 👎.

@github-actions
Copy link
Copy Markdown

github-actions bot commented Mar 25, 2026

Binary size checks passed

7 passed

Artifact Platform Gzip Baseline Delta Status
bridge_cffi Linux 3.9 MB 5.7 MB -1.8 MB (-30.9%) OK
bridge_cffi-stripped Linux 2.4 MB 4.3 MB -1.9 MB (-43.8%) OK
bridge_cffi macOS 3.2 MB 4.6 MB -1.4 MB (-30.6%) OK
bridge_cffi-stripped macOS 1.9 MB 3.5 MB -1.5 MB (-44.1%) OK
bridge_cffi Windows 3.2 MB 4.6 MB -1.4 MB (-31.4%) OK
bridge_cffi-stripped Windows 2.0 MB 3.5 MB -1.6 MB (-44.1%) OK
bridge_wasm WASM 1.9 MB 3.0 MB -1.0 MB (-35.0%) OK

Generated by cargo size-gate · workflow run

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
baml_language/crates/bex_engine/tests/orchestration.rs (1)

101-129: 🧹 Nitpick | 🔵 Trivial

Remove the ignored debug probe or make it assert on FB.

diag_fallback_fields never runs, has no assertions, and check() returns [A, B] instead of exercising FB, so it does not inspect fallback fields at all. This reads like leftover debugging code rather than regression coverage.


ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 7fcc270b-cb1e-4061-b5f4-edcb5669d19a

📥 Commits

Reviewing files that changed from the base of the PR and between 422ef5e and b11e050.

📒 Files selected for processing (4)
  • baml_language/crates/baml_builtins2/baml_std/baml/llm.baml
  • baml_language/crates/baml_builtins2/baml_std/baml/llm_types.baml
  • baml_language/crates/baml_compiler2_tir/src/builder.rs
  • baml_language/crates/bex_engine/tests/orchestration.rs

Comment on lines +35 to +47
function advance_round_robin(self) -> void {
match (self.client_type) {
ClientType.Primitive => {},
ClientType.Fallback => {
for (let sub in self.sub_clients) {
sub.advance_round_robin();
}
},
ClientType.RoundRobin => {
self.counter += 1;
},
}
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Advance the selected round-robin child too.

build_plan_with_state() consumes the chosen child subtree after computing idx, so nested round-robin nodes under that child can participate in the plan. advance_round_robin() only increments self.counter, though, so those nested counters never advance and the same inner branch is reused whenever this outer RR picks that child again. Mirror the same child selection here and recurse into the chosen sub-client as well.

Comment on lines +64 to +70
let result: T = execute_step(step, context) catch (e) {
_ => { continue; }
};
return result;
}

throw root.errors.DevOther { message: "All orchestration steps failed" };
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Don't treat every execute_step() failure as retryable.

execute_step() can fail with deterministic configuration errors too, not just provider failures. Catching _ here retries the next step and eventually replaces the real cause with "All orchestration steps failed", which can also send extra requests for a non-retryable bug. Restrict the fallback path to step-local provider/HTTP failures and rethrow the rest.

Comment on lines +2261 to +2285
Definition::Let(let_loc) => {
// Determine type from the let-binding's origin.
let db = self.context.db();
let item_tree =
baml_compiler2_hir::file_item_tree(db, let_loc.file(db));
let let_data = &item_tree[let_loc.id(db)];
match let_data.origin {
baml_compiler2_ast::ast::LetOrigin::Client => {
// client<llm> declarations produce Client instances.
Ty::Class(crate::ty::QualifiedTypeName::new(
baml_base::Name::new("baml"),
vec![baml_base::Name::new("llm")],
baml_base::Name::new("Client"),
))
}
baml_compiler2_ast::ast::LetOrigin::RetryPolicy => {
// retry_policy declarations produce RetryPolicy instances.
Ty::Class(crate::ty::QualifiedTypeName::new(
baml_base::Name::new("baml"),
vec![baml_base::Name::new("llm")],
baml_base::Name::new("RetryPolicy"),
))
}
_ => Ty::Unknown,
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Handle let-bound globals in package-qualified paths too.

This branch only fixes bare identifiers through infer_single_name(). root.A.build_plan() and package-qualified client/retry-policy references still go through resolve_package_item(), which only recognizes Definition::Function, so they continue to infer as Ty::Unknown and member lookup fails. Please share the same Definition::Let → type mapping with the multi-segment/package path resolver.

@sxlijin sxlijin changed the title baml-lanugage/fix: use build_plan for llm functions, delete execute baml-language/fix: use build_plan for llm functions, delete execute Mar 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant